from numpy import tile, array
import operator

group = array([
    [3, 104], [2, 100], [1, 81],
    [101, 10], [99, 5], [98, 2],
])
label = ["Romance"] * 3 + ["Action"] * 3

class KNNModel:
    r"""
    This Class generates kNN Models.
    @param dataSet {Numpy.array} training data.
    @param labels {array} labels of training data.
    @method predict(self {KNNModel}, inX {array}, k {int}) predicts the label of inX.
    """
    def __init__(self, dataSet, labels):
        self.dataSetSize = dataSet.shape[0]
        self.dataSet = dataSet
        self.labels = labels
    def predict(self, inX, k):
        diffMat = tile(inX, (self.dataSetSize, 1)) - self.dataSet
        sqDiffMat = diffMat ** 2
        sqDistances = sqDiffMat.sum(axis = 1)
        distances = sqDistances ** 0.5
        sortedDistIndicies = distances.argsort()
        classCount = {}
        for i in range(k):
            voteIlabel = self.labels[sortedDistIndicies[i]]
            classCount[voteIlabel] = classCount.get(voteIlabel, 0) + 1
        sortedClassCount = sorted(classCount.items(), key = operator.itemgetter(1), reverse = True)
        return sortedClassCount[0][0]

model = KNNModel(group, label)
print(model.predict([18, 90], 3))
print(model.predict([8, 49], 3))
print(model.predict([0, 1], 3))
print(model.predict([62, 62], 3))
print(model.predict([77, 1], 3))
print(model.predict([17, 0], 3))
